#include "ring_sampler.h"

/*
 * @brief 生成一个在 [0, mask] 范围内的随机整数。
 *
 * 该函数通过从 PRNG 中读取随机字节，并将其转换为 64 位整数，然后对结果进行掩码操作，以确保结果在 [0, mask] 范围内。
 *
 * @param prng 输入的 PRNG 对象。
 * @param mask 掩码值，范围0~255。
 * @return 生成的随机整数。
 */
uint64_t RandInt64(KeyedPRNG *prng, uint64_t mask)
{
    std::vector<unsigned char> randomBytes(8);
    prng->Clock(randomBytes);
    uint64_t randomUint64 = (static_cast<uint64_t>(randomBytes[0]) << 56 | static_cast<uint64_t>(randomBytes[1]) << 48 |
        static_cast<uint64_t>(randomBytes[2]) << 40 | static_cast<uint64_t>(randomBytes[3]) << 32 |
        static_cast<uint64_t>(randomBytes[4]) << 24 | static_cast<uint64_t>(randomBytes[5]) << 16 |
        static_cast<uint64_t>(randomBytes[6]) << 8 | static_cast<uint64_t>(randomBytes[7]));
    return mask & randomUint64;
}

BaseSampler::BaseSampler(KeyedPRNG *prng, Ring *baseRing) : prng(prng), baseRing(baseRing)
{
    // this->prng = new KeyedPRNG();
    // this->prng->clock = prng->clock;
}

BaseSampler::~BaseSampler() {}

UniformSampler::UniformSampler(KeyedPRNG *prng, Ring *baseRing) : BaseSampler(prng, baseRing)
{
    randomBufferN.resize(baseRing->N);
}

/*
 * @brief 从均匀分布中读取多项式。
 *
 * 该函数从均匀分布中读取多项式，并将其存储在 `Pol` 中。
 *
 * @param Pol 输入的多项式。
 */
void UniformSampler::Read(Poly *Pol)
{
    uint64_t randomUint, mask, qi;
    size_t ptr = 0;
    prng->Clock(randomBufferN);

#pragma omp parallel for
    for (size_t j = 0; j < baseRing->Modulus.size(); ++j) {
        qi = baseRing->Modulus[j];
        mask = baseRing->Mask[j];
        std::vector<uint64_t> &ptmp = Pol->Coeffs[j];
        for (size_t i = 0; i < baseRing->N; ++i) {
            do {
                if (ptr == baseRing->N) {
                    prng->Clock(randomBufferN);
                    ptr = 0;
                }
                randomUint = (static_cast<uint64_t>(randomBufferN[ptr]) << 56 |
                    static_cast<uint64_t>(randomBufferN[ptr + 1]) << 48 |
                    static_cast<uint64_t>(randomBufferN[ptr + 2]) << 40 |
                    static_cast<uint64_t>(randomBufferN[ptr + 3]) << 32 |
                    static_cast<uint64_t>(randomBufferN[ptr + 4]) << 24 |
                    static_cast<uint64_t>(randomBufferN[ptr + 5]) << 16 |
                    static_cast<uint64_t>(randomBufferN[ptr + 6]) << 8 |
                    static_cast<uint64_t>(randomBufferN[ptr + 7])) &
                    mask;
                ptr += 8;
            } while (randomUint >= qi);
            ptmp[i] = randomUint;
        }
    }
}

Poly *UniformSampler::ReadNew()
{
    Poly *Pol = baseRing->NewPoly();
    Read(Pol);
    return Pol;
}

/*
 * @brief 从均匀分布中读取多项式的特定级数。
 *
 * 该函数从均匀分布中读取多项式的特定级数，并将其存储在 `Pol` 中。
 *
 * @param level 要读取的级数。
 * @param Pol 输入的多项式。
 */
void UniformSampler::ReadLvl(int level, Poly *Pol)
{
    uint64_t randomUint, mask, qi;
    size_t ptr = 0;
    prng->Clock(randomBufferN);

    for (int j = 0; j < level + 1; j++) {
        qi = baseRing->Modulus[j];
        // Start by computing the mask
        mask = baseRing->Mask[j];
        std::vector<uint64_t> &ptmp = Pol->Coeffs[j];
        // Iterate for each modulus over each coefficient
        for (uint64_t i = 0; i < baseRing->N; i++) {
            // Sample an integer between [0, qi-1]
            while (true) {
                // Refill the pool if it runs empty
                if (ptr == baseRing->N) {
                    prng->Clock(randomBufferN);
                    ptr = 0;
                }
                // Read bytes from the pool
                randomUint = (static_cast<uint64_t>(randomBufferN[ptr]) << 56 |
                    static_cast<uint64_t>(randomBufferN[ptr + 1]) << 48 |
                    static_cast<uint64_t>(randomBufferN[ptr + 2]) << 40 |
                    static_cast<uint64_t>(randomBufferN[ptr + 3]) << 32 |
                    static_cast<uint64_t>(randomBufferN[ptr + 4]) << 24 |
                    static_cast<uint64_t>(randomBufferN[ptr + 5]) << 16 |
                    static_cast<uint64_t>(randomBufferN[ptr + 6]) << 8 |
                    static_cast<uint64_t>(randomBufferN[ptr + 7])) &
                    mask;
                ptr += 8;
                // If the integer is between [0, qi-1], break the loop
                if (randomUint < qi) {
                    break;
                }
            }
            ptmp[i] = randomUint;
        }
    }
}

TernarySampler::TernarySampler(KeyedPRNG *prng, Ring *baseRing, double p) : BaseSampler(prng, baseRing), p(p)
{
    initializeMatrix();
    if (p != 0.5) {
        computeMatrixTernary(p);
    }
}

void TernarySampler::Read(Poly *pol)
{
    sampleProba(baseRing->Modulus.size() - 1, pol);
}

void TernarySampler::initializeMatrix()
{
    matrixValues.resize(baseRing->Modulus.size());
    for (size_t i = 0; i < baseRing->Modulus.size(); i++) {
        uint64_t Qi = baseRing->Modulus[i];
        matrixValues[i][0] = 0;
        matrixValues[i][1] = 1;
        matrixValues[i][2] = Qi - 1;
    }
}

/*
 * @brief 计算三元多项式的概率矩阵。
 *
 * 该函数计算三元多项式的概率矩阵，并将其存储在 `matrixProba` 中。
 *
 * @param p 三元多项式的概率。
 */
void TernarySampler::computeMatrixTernary(double pCom)
{
    double g = pCom * std::pow(2.0, precision);
    uint64_t x = static_cast<uint64_t>(g);
    for (uint64_t j = 0; j < precision - 1; j++) {
        matrixProba[0][j] = static_cast<uint8_t>((x >> (precision - j - 1)) & 1);
    }
    g = (1 - pCom) * std::pow(2.0, precision);
    x = static_cast<uint64_t>(g);
    for (uint64_t j = 0; j < precision - 1; j++) {
        matrixProba[1][j] = static_cast<uint8_t>((x >> (precision - j - 1)) & 1);
    }
}

/*
 * @brief 从三元多项式中采样。
 *
 * 该函数从三元多项式中采样，并将其存储在 `pol` 中。
 *
 * @param lvl 要采样的级数。
 * @param pol 输入的多项式。
 */
void TernarySampler::sampleProba(uint64_t lvl, Poly *pol)
{
    if (std::abs(p - 0.5) <= 0.0001) {
        std::vector<uint8_t> randomBytesCoeffs(baseRing->N >> 3);
        std::vector<uint8_t> randomBytesSign(baseRing->N >> 3);
        prng->Clock(randomBytesCoeffs);
        prng->Clock(randomBytesSign);
        for (uint64_t i = 0; i < baseRing->N; i++) {
            uint64_t coeff = uint64_t(uint8_t(randomBytesCoeffs[i >> 3]) >> (i & 7)) & 1;
            uint64_t sign = uint64_t(uint8_t(randomBytesSign[i >> 3]) >> (i & 7)) & 1;
            uint64_t index = (coeff & (sign ^ 1)) | ((sign & coeff) << 1);
            for (int j = 0; j < lvl + 1; j++) {
                pol->Coeffs[j][i] = matrixValues[j][index];
            }
        }
    } else {
        std::vector<uint8_t> randomBytes(baseRing->N);
        uint8_t pointer = 0;
        uint64_t bytePointer = 0;
        prng->Clock(randomBytes);
        for (uint64_t i = 0; i < baseRing->N; i++) {
            auto [coeff, sign, updatedRandomBytes, updatedPointer, updatedBytePointer] =
                kysampling(prng, randomBytes, pointer, bytePointer, baseRing->N);
            pointer = updatedPointer;
            bytePointer = updatedBytePointer;
            uint64_t index = (coeff & (sign ^ 1)) | ((sign & coeff) << 1);
            for (uint64_t j = 0; j < lvl + 1; j++) {
                pol->Coeffs[j][i] = matrixValues[j][index];
            }
        }
    }
}
/*
 * @brief Knuth-Yao采样算法的实现。
 *
 * 该函数实现了Knuth-Yao离散分布采样算法,用于从给定的概率分布中采样。
 * 算法通过遍历概率矩阵的每一列,使用随机比特流来决定采样值。
 *
 * @param prng PRNG对象,用于生成随机比特。
 * @param randomBytes 存储随机字节的向量。
 * @param pointer 当前字节内的比特位置。
 * @param bytePointer 当前处理的字节位置。
 * @param byteLength 随机字节向量的长度。
 * @return 返回一个元组,包含:
 * - 采样得到的系数值
 * - 符号位
 * - 更新后的随机字节向量
 * - 更新后的比特位置
 * - 更新后的字节位置
 */

std::tuple<uint64_t, uint64_t, std::vector<uint8_t>, uint8_t, uint64_t> TernarySampler::kysampling(KeyedPRNG *prng,
    std::vector<uint8_t> &randomBytes, uint8_t pointer, uint64_t bytePointer, uint64_t byteLength)
{
    uint64_t sign = 0;
    int d = 0;
    int col = 0;
    int colLen = matrixProba.size();
    while (true) {
        for (int i = pointer; i < 8; i++) {
            d = (d << 1) + 1 - int((uint8_t(randomBytes[bytePointer]) >> i) & 1);
            if (d > colLen - 1) {
                return kysampling(prng, randomBytes, i, bytePointer, byteLength);
            }
            for (int row = colLen - 1; row >= 0; row--) {
                d -= int(matrixProba[row][col]);
                if (d == -1) {
                    if (i == 7) {
                        pointer = 0;
                        bytePointer++;
                        if (bytePointer >= byteLength) {
                            bytePointer = 0;
                            prng->Clock(randomBytes);
                        }
                        sign = uint8_t(randomBytes[bytePointer]) & 1;
                    } else {
                        pointer = i;
                        sign = uint8_t(randomBytes[bytePointer] >> (i + 1)) & 1;
                    }
                    return { uint64_t(row), uint64_t(sign), randomBytes, pointer + 1, bytePointer };
                }
            }
            col++;
        }
        pointer = 0;
        bytePointer++;
        if (bytePointer >= byteLength) {
            bytePointer = 0;
            prng->Clock(randomBytes);
        }
    }
}

void TernarySampler::ReadLvl(int lvl, Poly *pol)
{
    sampleProba(lvl, pol);
}

/*
 * @brief 高斯采样器的构造函数。
 *
 * 该函数初始化高斯采样器，设置随机数生成器、基环、标准差和边界。
 *
 * @param prng 输入的随机数生成器。
 * @param baseRing 输入的基环。
 * @param sigma 高斯分布的标准差。
 * @param bound 高斯分布的边界。
 */
GaussianSampler::GaussianSampler(KeyedPRNG *prng, Ring *baseRing, double sigma, uint64_t bound)
    : BaseSampler(prng, baseRing), sigma(sigma), bound(bound)
{
    randomBufferN.resize(1024);
    ptr = 0;
}

GaussianSampler::~GaussianSampler() {}

void GaussianSampler::Read(Poly *pol)
{
    ReadLvl(uint64_t(baseRing->Modulus.size() - 1), pol);
}

/*
 * @brief 从随机源读取高斯采样并将其添加到多项式的指定级别。
 *
 * 该函数从随机源读取高斯分布的系数，并将其添加到指定级别的多项式的系数中。
 * 采样的系数通过将浮动系数转化为整数并进行缩放，最终结果根据符号更新多项式的系数。
 *
 * @param level 目标级别，指定了多项式中要更新的级别。
 * @param pol 目标多项式，采样的系数将会添加到该多项式的系数中。
 */
void GaussianSampler::ReadLvl(uint64_t level, Poly *pol)
{
    double coeffFloat;
    uint64_t coeffInt, sign;
    prng->Clock(randomBufferN);
    for (uint64_t i = 0; i < baseRing->N; i++) {
        do {
            normFloat64(coeffFloat, sign);
            coeffInt = static_cast<uint64_t>(coeffFloat * sigma + 0.5);
        } while (coeffInt > bound);
        for (uint64_t j = 0; j <= level; j++) {
            uint64_t mod = baseRing->Modulus[j];
            pol->Coeffs[j][i] = (coeffInt * sign) | ((mod - coeffInt) * (sign ^ 1));
        }
    }
}

/*
 * @brief 从随机源读取高斯采样并将其添加到多项式的指定级别。
 *
 * 该函数从随机源读取高斯分布的系数，并将其与指定级别的多项式的系数相加。
 * 采样的系数通过对其进行放大和规范化后与模数相加，最终结果对每个级别的模数进行模运算。
 *
 * @param level 目标级别，指定了多项式中要更新的级别。
 * @param pol 目标多项式，采样的系数将会添加到该多项式的系数中。
 */
void GaussianSampler::ReadAndAddLvl(uint64_t level, Poly *pol)
{
    double coeffFloat;
    uint64_t coeffInt, sign;
    prng->Clock(randomBufferN);
    for (uint64_t i = 0; i < baseRing->N; i++) {
        do {
            normFloat64(coeffFloat, sign);
            coeffInt = static_cast<uint64_t>(coeffFloat * sigma);
        } while (coeffInt > bound);
        for (uint64_t j = 0; j <= level; j++) {
            uint64_t mod = baseRing->Modulus[j];
            pol->Coeffs[j][i] = CRed(pol->Coeffs[j][i] + ((coeffInt * sign) | ((mod - coeffInt) * (sign ^ 1))), mod);
        }
    }
}

/*
 * @brief 从随机字节序列中生成一个 64 位浮动值。
 *
 * 该函数将一个字节数组中的 8 个字节转换为一个 64 位无符号整数，并从该整数中提取出有效的浮动值。
 *
 * @param randomBytes 输入的字节序列，包含生成随机数所需的字节数据。
 * @param ptr 当前字节序列的偏移位置，指示从哪个位置开始读取字节。
 *
 * @return 生成的浮动值，范围在 [0, 1) 之间。
 */
double GaussianSampler::randFloat64(std::vector<uint8_t> &randomBytes, uint64_t ptr)
{
    uint64_t tmp = (static_cast<uint64_t>(randomBytes[0 + ptr]) << 56) |
        (static_cast<uint64_t>(randomBytes[1 + ptr]) << 48) | (static_cast<uint64_t>(randomBytes[2 + ptr]) << 40) |
        (static_cast<uint64_t>(randomBytes[3 + ptr]) << 32) | (static_cast<uint64_t>(randomBytes[4 + ptr]) << 24) |
        (static_cast<uint64_t>(randomBytes[5 + ptr]) << 16) | (static_cast<uint64_t>(randomBytes[6 + ptr]) << 8) |
        (static_cast<uint64_t>(randomBytes[7 + ptr]));
    return double(tmp & 0x1fffffffffffff) / double(0x1fffffffffffff);
}
/*
 * @brief 生成一个高斯分布的浮动值。
 *
 * 该函数生成符合高斯分布的浮动系数，并返回其值及符号。通过基于随机字节生成的均匀分布数，经过几次数值变换，最终生成高斯分布的数值。
 * 函数采用了 Box-Muller 算法来生成符合标准正态分布的随机数。
 *
 * @param coeffFloat 输出的高斯分布浮动系数。
 * @param resSign 输出的符号，值为 1 或 -1，表示采样值的符号。
 */
void GaussianSampler::normFloat64(double &coeffFloat, uint64_t &resSign)
{
    double x, y;
    uint64_t sign;
    int32_t j;
    while (true) {
        if (ptr == randomBufferN.size()) {
            prng->Clock(randomBufferN);
            ptr = 0;
        }
        uint32_t juint32 = (randomBufferN[ptr] << 24) | (randomBufferN[ptr + 1] << 16) | (randomBufferN[ptr + 2] << 8) |
            randomBufferN[ptr + 3];
        ptr += 8;
        j = int32_t(juint32 & 0x7FFFFFFF);
        sign = juint32 >> 31;
        int32_t i = j & 0x7F;
        x = static_cast<double>(j) * wn[i];
        if (uint32_t(j) < kn[i]) {
            coeffFloat = x, resSign = sign;
            return;
        }
        if (i == 0) {
            while (true) {
                if (ptr == uint64_t(randomBufferN.size())) {
                    prng->Clock(randomBufferN);
                    ptr = 0;
                }

                x = -log(randFloat64(randomBufferN, ptr)) * (1.0 / 3.442619855899);
                ptr += 8;

                if (ptr == uint64_t(randomBufferN.size())) {
                    prng->Clock(randomBufferN);
                    ptr = 0;
                }

                y = -log(randFloat64(randomBufferN, ptr));
                ptr += 8;

                if (y + y >= x * x) {
                    break;
                }
            }

            coeffFloat = x + 3.442619855899, resSign = sign;
            return;
        }
        if (ptr == uint64_t(randomBufferN.size())) {
            prng->Clock(randomBufferN);
            ptr = 0;
        }

        // 3
        float tmp1 = fn[i] + float(randFloat64(randomBufferN, ptr)) * (fn[i - 1] - fn[i]);
        float tmp2 = float(exp(-0.5 * x * x));
        if (tmp1 < tmp2) {
            ptr += 8;
            coeffFloat = x, resSign = sign;
            return;
        }
        ptr += 8;
    }
}